# $Id: libomd.py 14 2006-08-06 12:08:22Z niels $
#
# Copyright (C) 2006  Niels de Vos
#
# Modified by David Were
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
#

import socket
import cPickle
import thread
from threading import Semaphore, Thread, Lock
import time
import logging
import os
import sys
import getopt
import signal
import re

# default destination
defAddr = '<broadcast>'
defPort = 5000

# max size of a message
SOCK_BUF_SIZE = 2048

# internals
__log = None
__stdSock = None
__bcastSock = None
__recvQueue = list()
__listening = True
__stdRecvThread = None
__bcastRecvThread = None
# using a sem as blocking counter
__recvQueueSem = Semaphore(0)

PROC_STAT = [
    'pid',        'comm',        'state',    'ppid',
    'pgrp',        'session',    'tty_nr',    'tpgid',
    'flags',    'minflt',    'cminflt',    'majflt',
    'cmajflt',    'utime',    'stime',    'cutime',
    'cstime',    'priority',    'nice',        'num_threads',
    'itrealvalue',    'starttime',    'vsize',    'rss',
    'rlim',        'startcode',     'endcode',     'startstack',
    'kstkesp',    'kstkeip',    'signal',    'blocked',
    'sigignore',    'sigcatch',    'wchan',    'nswap',
    'cnswap',    'exit_signal',    'processor',    'rt_priority',
    'policy']


    
class ProcessInfo:
    """ A simple ProcessInfo class.
    This class contains meta-data about a process.
    """
    stat = None # dict()
    """ A dict() read from /proc/<pid>/stat, see 'man 5 proc'. """
    
    stayReasons = None
    """ Stayreasons from /proc/<pid>/om/stay in a list. """
    
    isMigrated = False
    """ Boolean for marking a process as migrated. """
    
    isBlacklisted = False
    """ Boolean for marking a process as blacklisted. """
    
            
class ProcessList:
    __procs = dict()
    """ List with processes (pid, ProcessInfo). """
    __lock = Lock()
    """ A lock for preventing access to the dict while in use. """
    __log = None
    """ The logging for this ProcessList. """
    __isRunning = True
    """ irRunning flag, used for stopping the process. """
    __updateThread = None
    """ Thread as created with thread.start_new_thread(). """
    
    def __init__(self, timeout = 5):
        self.__log = logging.getLogger('processlist')
        if not self.__log.handlers:
            self.__log.addHandler(logging.StreamHandler(sys.stderr))
        self.__updateThread = thread.start_new_thread(self.__updateList, (timeout, ))
        self.__log.info('started new thread')
        
    def stop(self):
        self.__isRunning = False
        self.__updateThread.join()
        
    def __updateList(self, _timeout):
        try:
            while self.__isRunning:
                for pid in self.__getPids():
                    self.__lock.acquire()
                    proc = self.__getProcessInfo(pid)
                    self.__lock.release()
                    if not proc:
                        self.__log.debug('Process %s finished while getting stats' % pid)
                time.sleep(_timeout)
        except Exception, e:
            self.__log.critical('BUG in ProcessList.__updateList: %s' % e)
    
    def __getProcessInfo(self, _pid):
        """ Read the ProcessInfo from the proc-filesystem.
        When using this method, the __lock should be acquired.
        """
        proc = ProcessInfo()
        proc.stat = readPid(_pid)
        if not proc.stat:
            return None
        if not self.__procs.has_key(_pid) or self.__procs[_pid].stat['starttime'] != proc.stat['starttime']:
            self.__procs[_pid] = proc
            proc.stayReasons = self.__getStayReasons(_pid)
            if proc.stayReasons:
                proc.isBlacklisted = True
        return proc
    
    def __getStayReasons(self, _pid):
        """ read /proc/<pid>/om/stay and return as a list """
        stayFile = '/proc/%s/om/stay' % _pid
        stayReasons = []
        f = None
        try:
            f = file(stayFile)
            stayReasons = f.readlines()
            stayReasons = map(lambda i: i.strip(), stayReasons)
        except IOError, ie:
            # could not read process
            # reason: not owner or process finished
            self.__log.warning('could not read %s' % stayFile)
            stayReasons = True
        if f:
            f.close()
        return stayReasons
    
    def getProcessInfo(self, _pid):
        """ Return the ProcessInfo for a pid. """
        proc = None
        self.__lock.acquire()
        if not self.__procs.has_key(_pid):
            proc = self.__getProcessInfo(_pid)
        else:
            proc = self.__procs[_pid]
        self.__lock.release()
        return proc
    
    def __getPids(self):
        """ get all processes from /proc """
        pids = os.listdir('/proc')
        pids = filter(lambda i: re.match('^\d+$', i), pids)
        # convert pids from strings to integers
        pids = map(int, pids)
        # remove all processes from __procs who are not in pids
        self.__lock.acquire()
        for pid in self.__procs.keys():
            if pid not in pids:
                self.__procs.pop(pid)
        self.__lock.release()
        return pids

    def getPids(self):
        """ Return all pids in the list. """
        return self.__procs.keys()


class nodeInfo:
    """a Flag indicating if this is the home node"""
    home = False
    """information on cpus"""
    cpuinfo = None
    
    """stores current cpu load"""
    load = None
    
    """stores current memrory usage"""
    meminfo = dict()
    
    """stores time stamp of last status update"""
    lastUpdate = float()
    
    ip =''

class nodeList:
    """ A Dictionary containing (ip,nodeinfo)"""
    __nodes = dict()
    
    def init(self):
        self.__log = logging.getLogger('nodelist')
        if not self.__log.handlers:
            self.__log.addHandler(logging.StreamHandler(sys.stderr))
            
    def addNode(self,src,cpuinfo,meminfo,home=False):
        """create new nodeInfo instance to store information"""
        node = nodeInfo()
        
        """Populate NodeInfo class with Information"""
        node.cpuinfo = cpuinfo
        node.meminfo = meminfo
        node.lastUpdate = time.time()
        node.home = home
        """Add nodeInfo to Dictionary"""
        self.__nodes[src] = node
    
    def updateNode(self,src,load,meminfo):
        if self.contains(src):
            self.__nodes[src].load = load
            self.__nodes[src].meminfo = meminfo
            self.__nodes[src].lastUpdate = time.time()

    def contains(self,src):
        return self.__nodes.has_key(src)
    
    def __updateList(self):
        if self.__nodes.keys():
            for x in self.__nodes.keys():
                    if time.time() - self.__nodes[x].lastUpdate > 15:
                        del self.__nodes[x]
                
    def getNodes(self):
        self.__updateList()
        return self.__nodes
        
        
        
def omCheck():
    """ Check if we're running on OM.
        currently simply checks if /proc/hpc exists.
        NOTE: maybe use some unique version here?
    """
    return os.path.exists('/proc/hpc') 
    

def initLogging():
    global __log
    if not __log:
        __log = logging.getLogger('libomd')
        if not __log.handlers:
            __log.addHandler(logging.StreamHandler(sys.stderr))
    
def initNetworking(__bind = '', __recvBcast = False):
    """ Networking setup.
        Start at least one thread for receiving. A second thread when
        receiving on '<breadcast>' too.
    """
    global __stdSock, __recvSock, __listening, __bcastRecvThread, __log
    
    # setup the standard socket (receiving & sending)
    __stdSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    __stdSock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    __stdRecvThread = thread.start_new_thread(__recvLoop, (__stdSock, __recvQueueSem))
    try:
        __stdSock.bind((__bind, defPort))
    except socket.error, e:
        __log.error('Error: %s' % e)

    # setup a broadcast socket (receiving)
    if __bind != '' and __recvBcast:
        __bcastSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        __bcastSock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        try:
            __bcastSock.bind(('<broadcast>', defPort))
            __bcastRecvThread = thread.start_new_thread(__recvLoop, (__bcastSock, __recvQueueSem))
        except socket.error, e:
            __log.error('Error: %s' % e)
            __log.error('- still running, but unable to receive broadcast messages')

def stopNetworking():
    """ Cleanup networking.
        shutwdown sockets nicely etc.
    """
    global __listening, __stdSock, __bcastSock, __log
    __listening = False
    try:
        __stdSock.close()
    except socket.error, e:
        __log.error('stopNetworking: could not shutdown stdSock...')
        __log.debug('stopNetworking: %s' % e)
    if __bcastSock:
        try:
            __bcastSock.close()
        except socket.error, e:
            __log.error('stopNetworking: could not shutdown bcastSock...')
            __log.debug('stopNetworking: %s' % e)
    
def sendMsg(msg):
    """ Send a message to 'msg.dst' or 'defAddr'.
        The message is converted using the cPickle module.
    """
    global __stdSock, __log
    rawMsg = cPickle.dumps(msg)
    __log.debug('sendMsg: %s' % msg)
    dst = defAddr
    if msg.dst != None:
        dst = msg.dst
    __log.debug('sendMsg: sending %d bytes to %s:%d' % (len(rawMsg), dst, defPort))
    try:
        size = __stdSock.sendto(rawMsg, (dst, defPort))
        if size != len(rawMsg):
            __log.debug('sendMsg: tried to send %d bytes, %d sent...' % (len(rawMsg), size))
    except socket.error, e:
        __log.error('sendMsg: Failed to send message...')
        __log.error('sendMsg: Error: %s' % e)

def recvMsg():
    """ A blocking receive function.
        Wait for messages in the __recvQueue. When the queue is empty, block until
        one of the receive threads releases the blocking Semaphore.
    """
    global __recvQueue, __stdSock, __recvQueueSem, __log
    msg = None
    __log.debug('recvMsg(): waiting for message...')
    while not msg:
        if len(__recvQueue) > 0:
            msg =  __recvQueue.pop(0)
        else:
            __recvQueueSem.acquire()        
    __log.debug('recvMsg(): returning %s - %d msgs left' % (msg, len(__recvQueue)))
    return msg

def __sockRecvMsg(__sock):
    """ Function used internally for receiving a msg on a socket.
        This function is used by the standard receive thread and when
        running, the broadcast receive thread.
        Converting raw data to Pythong Objects is handled by cPickle.
    """
    global __log
    src = None
    msg = None
    rawMsg = None
    try:
        msgTuple = __sock.recvfrom(SOCK_BUF_SIZE)
        if msgTuple:
            (rawMsg, addr) = msgTuple
            if addr:
                (src, port) = addr
    except socket.timeout, t:
        pass
    if rawMsg:
        try:
            msg = cPickle.loads(rawMsg)
            msg.src = src
        except ValueError, ve:
            __log.error('Could not unpickle:\n%s\n' % rawMsg)
        except cPickle.UnpicklingError, pe:
            __log.error('__sockRecvMsg: could not read message: %s' % pe)
    return msg

def __recvLoop(__sock, __sem):
    """ Thread function for receiving messages.
        Started by the standard receive thread and when running, the
        broadcast receive thread.
    """
    global __recvQueue, __listening, __log
    __log.debug('__recvLoop(): started...')
    while __listening:
        msg = __sockRecvMsg(__sock)
        __log.debug('__recvLoop(): received %s' % msg)
        __recvQueue.append(msg)
        __sem.release()

def migrate(_pid, _dst):
    """ Start migration of a _pid to _dst (ip-address). """
    global __log
    try:
        os.system('echo '+_dst+' >/proc/%s/om/where' % _pid)
    except IOError, ie:
        __log.error('Process %s can not be migrated...' % _pid)
    __log.debug('migrated %s to %s' % (_pid, _dst))
    

def readPid(_pid):
    global __log
    proc = None
    if os.path.exists('/proc/%s/stat' % _pid):
        try:
            stat = file('/proc/%s/stat' % _pid, 'r')
            proc = dict(zip(PROC_STAT, re.split('\s', stat.readline())))
        except OSError, e:
                __log.error('Error reading /proc/stat:\n%s\n' % e)
    return proc

class Omd:
    isRunning = True
    log = None
    broadcast = False
    # timeout in secs between status messages
    delay = 5
    processList = None

    def __init__(self):
        if not self.log:
            self.log = logging.getLogger('omd')
            if not self.log.handlers:
                self.log.addHandler(logging.StreamHandler(sys.stderr))
        
    def mainLoop(self, delay):
        self.log.critical('Programming error: override Omd.mainLoop()')
    
    def recvLoop(self):
        """ A thread which receives and handles messages. """
        while self.isRunning:
            msg = recvMsg()
            self.log.debug('recvMsg: received message - %s' % msg)
            self.handleMsg(msg)
        self.log.debug('recvLoop: exiting...')
    
    def handleMsg(self, msg):
        self.log.critical('Programming error: override Omd.handleMsg()')
    
    def start(self, listen = ''):
        ret = 0
        
        if omCheck():
            # setup logging
            initLogging()
            # setup networking
            initNetworking(listen, self.broadcast)
    
            # start threads and wait for them to finish
            #statusThread = thread.start_new_thread(sendStatusLoop, (defTimeout, ))
            mainThread = thread.start_new_thread(self.mainLoop, (self.delay, ))
            recvThread = thread.start_new_thread(self.recvLoop, ())
            try:
                signal.pause()
            except (KeyboardInterrupt, int):
                pass
            except (Exception, sig):
                self.log.debug('Caugth signal: %s' % sig)
            self.isRunning = False
            if self.processList:
                self.processList.stop()
            stopNetworking()
        else:
            self.log.error('You\'re not running on OpenMosix, exiting...')
            ret = 1
        sys.exit(ret)
